import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from scipy.interpolate import griddata
import matplotlib.colors as colors
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.patches import Rectangle
import io
import json

temp = []
IPValue = {}
fig, ax = plt.subplots(figsize=(15, 7.5))


def ERTuiArray(ABMN, a1_dtype, a1_shape, Aspacing):
    abmn = ABMN.encode('ISO-8859-1')
    ABMN1 = np.fromstring(abmn, dtype=getattr(np, a1_dtype)).reshape(eval(a1_shape))
    # print(print(ABMN1[0,:]))
    return ABMN1

#拟断面开始
def AnalyzeJonsJson(jb_tu,json_str, ns,nss):
    # 将JSON字符串转换为Python对象
    data = json.loads(json_str)
    # ReadData(data)
    io = ERTuiSection(data,jb_tu, int(ns), ns,nss)
    return io


def ERTuiSectionNs(json_str):
    data = json.loads(json_str)
    # print("json abmn =======================================")
    # print(data)
    #NS = np.sum(range(1,int(data['Max_n_spacing'])+1)) + int(data['Max_n_spacing'])*((int(data['num_electrodes'])-3)-int(data['Max_n_spacing']))      #总排列数
    NS = len(ERTuiArray(data['ABMN'], data['dtype'], data['shape'], data['Aspacing']))
    print(NS)
    return int(NS)

def ERTuiSection(data,jb_tu, AI, ns,nss):
    # NS = np.sum(range(1,int(data['Max_n_spacing'])+1)) + int(data['Max_n_spacing'])*((int(data['num_electrodes'])-3)-int(data['Max_n_spacing']))      #总排列数
    ###3. 读取CVS数据及系统参数
    '''
        手部首先需要向主机发送数据采集的命令，规定采集的参数，然后主机才能返回数据。这里假设命令发送成果，数据成功上传。
        数据按每次叠加上传，例如我们设置命令为叠加5次，每次波形长度为8，那么手部会收到8个完整波形后再做相应的处理。
        这8个波形的数据在数据文件中是一个数据块，有相应的头文件与之对应。
    '''
    Data = jb_tu
    ABMN = ERTuiArray(data['ABMN'], data['dtype'], data['shape'], data['Aspacing'])
    # list1 = list(map(int,data['ABMN']))
    #
    # ABMN = np.array(list1)
    # IPValue = {}
    global IPValue
    IPValue = IPValue;
    #plt.ion()  # 打开交互模式
    # for AI in range(NS):
    #Data = ReadData(data, ABMN[AI, :])
    print("----------------------ABMN[AI,:]------------")
    # print(ABMN)
    print(AI)
    data['ABMN'] = ABMN
    ###4. 处理数据

    IPValue[AI] = ComputeRhoIP(Data, data, AI)
    #print(IPValue)
    ###5. 画出图形
    io = ERTPlot(IPValue, ABMN, int(data['Aspacing']), ns,nss)
    return io

def ComputeRhoIP(Data, data, AI):
    '''
        计算电阻，极化率，及画图的位置.
        此算法暂时只针对偶极-偶极排列
        :param Data:
        :param SysInfo:
        :param SurveyInfo:
        :param ABMN:
        :return: IPValues
    '''
    # ABMN1 = ERTuiArray(data['ABMN'], data['dtype'], data['shape'], data['Aspacing'])
    # ABMN = ABMN1[AI, :]
    ABMN1 = data['ABMN']
    ABMN = ABMN1[AI, :]
    # print("----------------------------------------ABMN[AI,:]------------")
    # print(ABMN)
    SR = int(data['RXSR'])  # 数据点
    TXFrq = int(data['TXFrq'])  # 发射频率
    StackNO = int(data['StackNO'])  # 叠加次数  这三个都是app填写的
    Aspacing = int(data['Aspacing'])
    TXA = (int(float(ABMN[0])) - 1) * Aspacing  # Aspacing 点距 也是app 填写的
    TXB = (int(float(ABMN[1])) - 1) * Aspacing
    RXM = (int(float(ABMN[2])) - 1) * Aspacing
    RXN = (int(float(ABMN[3])) - 1) * Aspacing


    IPValues = {}
    IPValues['TXA'] = TXA
    IPValues['TXB'] = TXB
    IPValues['RXM'] = RXM
    IPValues['RXN'] = RXN
    IPValues['X'] = (RXM - TXB) / 2 + TXB  # 画图横坐标位置
    IPValues['Z'] = (RXM - TXB) / 2 + Aspacing / 2  # 画图垂向坐标位置
    IPValues['n'] = ABMN[2] - ABMN[1]  # 隔离系数

    print("--------------IPValues['n']---------------")
    print(IPValues['n'])

    #print(type(Data))
    #datas = Data.replace("'"," ");
    #dict = json.loads(s=Data)
    dict = eval(Data)
    # 使用字典推导式将字符串键转换为整数键
    Data = {int(key): value for key, value in dict.items()}
    #print(type(Data),Data)
    # 将字典中的值转换为浮点数类型
    for key in Data:
        # new_dict[key] = float(eval(new_dict[key]))
        Data[key] = np.array(Data[key])
   # print("-----------Data --------")
    #print(type(Data),Data)
    # 叠加处理数据
    temp =  Data[0]  # 将第一个叠加数据暂存起来
    # print(type(temp),temp)
    for BI in range(1, StackNO):
        #print(StackNO,BI)
        temp = temp + np.array(Data[BI])  # 叠加数据
    temp = temp / StackNO  # 平均数据
    #print(type(temp),temp)
    #print("---------------- SR ="+str(SR))
   #print("---------------- TXFrq ="+str(TXFrq))
    #NP = int(SR/TXFrq)  # 一个周期的数据点数
    NP = SR  # 一个周期的数据点数
    #print("---------------- np =")
   # print(NP)
    NPhlf = int(NP / 2)  # 半个周期数据点数
    #print("---------------- NPhlf ="+str(NPhlf))
    tmp = (temp[0:NPhlf, :] - temp[NPhlf:, :]) / 2  # 将后半个周期和前半个周期平均
    # print("---------------- tmp ="+str(tmp))
    TX_I = tmp[:, 0]  # 发射电流数据
    # print("---------------- TX_I ="+str(TX_I))
    RX_V = tmp[:, 1]  # 接收电压数据
    VpPnts = int(0.0625 * NP)  # 数据点数用于计算电流和电压
    TXI = np.sum(TX_I[int(NPhlf / 2 - VpPnts):int(NPhlf / 2)]) / VpPnts  # 发射电流值
    Vp = np.sum(RX_V[int(NPhlf / 2 - VpPnts):int(NPhlf / 2)]) / VpPnts  # 接收电压值
    IPValues['TXI'] = TXI
    IPValues['RXV'] = Vp
    Mpts = int(0.08125 * NP)
    Mst = int(0.05625 * NP) + 4
    IP = 1 / TXFrq / NP * 1.87 / Vp * np.sum(RX_V[int(NPhlf + Mst):int(NPhlf + Mst + Mpts)]) / 10
    IPValues['IP'] = IP
    print("----------Rho-----------")
    print(Vp,TXI,np.pi,IPValues['n'],Aspacing)
    Rho = 2.0/2.0 * np.pi * IPValues['n'] * (IPValues['n'] + 1) * (IPValues['n'] + 2) * Aspacing
    IPValues['Rho'] = Rho

    print(">>>>>>>>>>>>>>>>>>>>>IPValues= ")
    print(IPValues)
    return IPValues


def ERTPlot(IPValues, ABMN, Aspacing, NS,NSS):
   # print("___________________ERTPlot__________________11111")
    '''
        本程序画出高密度电阻率法的等值线图和排列图
        :param IPValues:
        :return: NONE
    '''
    ####这里画出两个图，一个是排列图，利用红色圆圈表示发射，用黑色圆表示接收
    # 最后一个IP的值
    # 获取相应的画图参数
    # print("-------------------------IPValues---------------------------")
    # print(IPValues)
    NIP = len(IPValues)
    print("-------------------------NIP---------------------------")
    print(NIP)
    lastIP = IPValues[NIP - 1]

    X = lastIP['X']
    Z = lastIP['Z']
    rho = lastIP['Rho']
    # #####计算图的尺寸
    #print(ABMN[:, 0])
    Min_X = np.min(ABMN[:, 0] - 1) * Aspacing
    Max_X = np.max(ABMN[:, 3] - 1) * Aspacing
    Min_Z = np.min(ABMN[:, 2] - ABMN[:, 1]) * Aspacing / 2 + Aspacing / 2
    Max_Z = np.max(ABMN[:, 2] - ABMN[:, 1]) * Aspacing / 2 + Aspacing / 2

    # fig1 = plt.figure(num=1)
    # 
    # ax = fig1.add_subplot(111)
    # ax.legend(frameon=False)

    colormap = cm.get_cmap('jet')
    rho_range = np.logspace(1, 5, 40)

    plt.xlim([Min_X, Max_X])
    plt.ylim([0, Max_Z + 10])
    ###画出所有的IPvalues
    rho_trans = (np.log10(rho) - np.log10(min(rho_range))) / 4
    cls = colormap(rho_trans)
    # 画为圆或者方块都可以，这里选择了方块，使得两个点之间的电阻率图像可以连接起来
    # Rho_C = plt.Circle((X, Z), radius=Aspacing / 3, fill=True, color=cls, linewidth=2.0, \
    #                   clip_on=False)   # 拟断面图画图点
    # plt.gca().add_artist(Rho_C)
    Rho_rect = Rectangle((X - Aspacing / 2, Z - Aspacing / 4), Aspacing, Aspacing / 2 * 0.95,
                         fill=True, color=cls, \
                         clip_on=False)  # 拟断面图画图点
    plt.gca().add_patch(Rho_rect)  # 绘制图像形状方形
    plt.annotate(format(rho, '.0f'), xy=(X, Z), fontsize=8, ha='center', va='center')
    plt.gca().invert_yaxis()
    plt.gca().set_xlabel('Array Location(m)', size=14)
    plt.gca().set_ylabel('n-spacing', size=14)
    plt.gca().xaxis.set_label_position('top')
    plt.gca().xaxis.tick_top()
    plt.gca().set_aspect('equal')
    plt.yticks(np.insert(np.arange(Min_Z, Max_Z + Aspacing / 2, Aspacing / 2), 0, 0))
    plt.xticks(np.arange(Min_X, Max_X, Aspacing))
    labels = [item.get_text() for item in ax.get_yticklabels()]
    NLAB = len(labels)
    labs = np.insert(np.arange(int(np.min(ABMN[:, 2] - ABMN[:, 1])), \
                               int(np.max(ABMN[:, 2] - ABMN[:, 1]) + 1), 1), 0, 0)
    print("-----------------------------------NLAB--------------------------------------------")
    print(NIP)
    print(NSS)
    if NIP == NSS:
        plt.clf()
        ax.axes.get_xaxis().set_visible(False)
        CX = np.zeros((NSS,))
        CZ = np.zeros((NSS,))
        Rho = np.zeros((NSS,))
        for I in range(NSS):
            currentIP = IPValues[I]
            CX[I] = currentIP['X']
            CZ[I] = currentIP['Z']
            Rho[I] = currentIP['Rho']
        #网格化
        xv, zv = np.meshgrid(np.linspace(Min_X, Max_X, 100), np.linspace(Min_Z, Max_Z, 100))
        #插值
        rho_av = griddata((CX, CZ), Rho, (xv, zv))
        cs = plt.contourf(xv, zv, rho_av, levels=np.logspace(1,5,40), cmap=cm.jet,norm = colors.LogNorm())
        cs.set_clim(vmin=10, vmax=10000)
        # plt.gca().clabel(cs, fmt='%2.1f', colors='k', fontsize=8, inline=True,inline_spacing=0)
        plt.gca().invert_yaxis()
        plt.gca().set_xlabel('Array Location(m)', size=14)
        plt.gca().set_ylabel('n-spacing', size=14)
        plt.gca().xaxis.set_label_position('top')
        plt.gca().xaxis.tick_top()
        plt.gca().set_aspect('equal')
        plt.yticks(np.insert(np.arange(Min_Z, Max_Z + Aspacing / 2, Aspacing / 2), 0, 0))
        plt.xticks(np.arange(Min_X, Max_X, Aspacing))
        labels = [item.get_text() for item in ax.get_yticklabels()]
        NLAB = len(labels)
        labs = np.insert(np.arange(int(np.min(ABMN[:, 2] - ABMN[:, 1])), \
                                   int(np.max(ABMN[:, 2] - ABMN[:, 1]) + 1), 1), 0, 0)
        if NLAB == len(labs):
            ax.set_yticklabels(labs)
        divider = make_axes_locatable(ax)
        cax = divider.append_axes('right', size="2%", pad=0.1)
        sm = plt.cm.ScalarMappable(norm=colors.LogNorm(), cmap=colormap)
        sm.set_clim(vmin=10, vmax=10000)
        cbar = plt.colorbar(sm, cax=cax)
        cbar.set_label('Apparent Resistivity(Ohm-m)', rotation=270, labelpad=20)
        ax.spines['bottom'].set_visible(False)
        plt.show()
        #cbar.remove()
        f = io.BytesIO()
        plt.savefig(f, format="png")
        print("___________________ERTPlot__________________33333")
        return f.getvalue()
    f = io.BytesIO()
    plt.savefig(f, format="png")
    print("___________________ERTPlot__________________22222")
    return f.getvalue()